#!/usr/bin/env python3
"""
从results目录下的restricted_sp_***MeV.csv文件中提取第7行数据
将数据存储到长度为5400的numpy数组中，缺失文件对应位置填0
"""

import os
import numpy as np
import pandas as pd
from pathlib import Path

def extract_seventh_row_data(results_dir="results1", output_file="ionization_data.npy"):
    """
    从CSV文件中提取第7行数据并保存为numpy数组
    
    Args:
        results_dir (str): 结果目录路径
        output_file (str): 输出的.npy文件名
    """
    
    # 初始化长度为5400的数组，默认值为0
    data_array = np.zeros(5400, dtype=float)
    
    # 统计信息
    found_files = 0
    missing_files = 0
    error_files = []
    
    print(f"开始处理 {results_dir} 目录下的CSV文件...")
    print(f"目标数组长度: {len(data_array)}")
    print("-" * 50)
    
    # 检查results目录是否存在
    if not os.path.exists(results_dir):
        print(f"错误: 目录 {results_dir} 不存在")
        return
    
    # 遍历1到5400
    for i in range(1, 5401):
        filename = f"ionization_{i}.csv"
        filepath = os.path.join(results_dir, filename)
        
        # 显示进度（每处理100个文件显示一次）
        if i % 100 == 0:
            print(f"处理进度: {i}/5400 ({i/5400*100:.1f}%)")
        
        if os.path.exists(filepath):
            try:
                # 读取CSV文件
                df = pd.read_csv(filepath, header=None)
                
                # 检查是否有足够的行数
                if len(df) >= 7:
                    # 获取第7行的数据（索引为6，因为从0开始）
                    seventh_row = df.iloc[6]  # 第7行
                    
                    # 如果第7行有多列，取第一列的值
                    # 如果需要取特定列，可以修改这里
                    if len(seventh_row) > 0:
                        value = seventh_row.iloc[0]  # 取第一列的值
                        
                        # 尝试转换为数字
                        try:
                            numeric_value = float(value)
                            data_array[i-1] = numeric_value  # 索引从0开始，所以i-1
                            found_files += 1
                        except (ValueError, TypeError):
                            print(f"警告: 文件 {filename} 第7行第1列无法转换为数字: {value}")
                            error_files.append(filename)
                    else:
                        print(f"警告: 文件 {filename} 第7行为空")
                        error_files.append(filename)
                else:
                    print(f"警告: 文件 {filename} 行数不足7行，实际行数: {len(df)}")
                    error_files.append(filename)
                    
            except Exception as e:
                print(f"错误: 读取文件 {filename} 时发生异常: {str(e)}")
                error_files.append(filename)
        else:
            missing_files += 1
            # 对应位置已经是0，不需要额外操作
    
    # 保存数组为.npy文件
    try:
        np.save(output_file, data_array)
        print(f"\n数据已成功保存到: {output_file}")
    except Exception as e:
        print(f"\n错误: 保存文件时发生异常: {str(e)}")
        return
    
    # 输出统计信息
    print("\n" + "="*60)
    print("处理完成统计:")
    print(f"成功读取的文件数量: {found_files}")
    print(f"缺失的文件数量: {missing_files}")
    print(f"读取错误的文件数量: {len(error_files)}")
    print(f"数组中非零值的数量: {np.count_nonzero(data_array)}")
    print(f"数组中零值的数量: {len(data_array) - np.count_nonzero(data_array)}")
    
    # 显示数组的基本统计信息
    non_zero_data = data_array[data_array != 0]
    if len(non_zero_data) > 0:
        print(f"\n非零数据统计:")
        print(f"最小值: {np.min(non_zero_data):.6f}")
        print(f"最大值: {np.max(non_zero_data):.6f}")
        print(f"平均值: {np.mean(non_zero_data):.6f}")
        print(f"标准差: {np.std(non_zero_data):.6f}")
    
    # 如果有错误文件，显示前几个作为示例
    if error_files:
        print(f"\n错误文件示例（前10个）:")
        for filename in error_files[:10]:
            print(f"  - {filename}")
        if len(error_files) > 10:
            print(f"  ... 还有 {len(error_files) - 10} 个错误文件")
    
    print("="*60)

def load_and_inspect_data(npy_file="extracted_data.npy"):
    """
    加载并检查保存的numpy数组
    
    Args:
        npy_file (str): .npy文件路径
    """
    try:
        data = np.load(npy_file)
        print(f"\n加载的数组信息:")
        print(f"数组形状: {data.shape}")
        print(f"数据类型: {data.dtype}")
        print(f"非零元素数量: {np.count_nonzero(data)}")
        print(f"前10个元素: {data[:10]}")
        print(f"后10个元素: {data[-10:]}")
        
        # 找到第一个和最后一个非零元素的位置
        non_zero_indices = np.nonzero(data)[0]
        if len(non_zero_indices) > 0:
            print(f"第一个非零元素位置: {non_zero_indices[0] + 1} (对应文件: restricted_sp_{non_zero_indices[0] + 1}MeV.csv)")
            print(f"最后一个非零元素位置: {non_zero_indices[-1] + 1} (对应文件: restricted_sp_{non_zero_indices[-1] + 1}MeV.csv)")
        
        return data
    except Exception as e:
        print(f"错误: 无法加载文件 {npy_file}: {str(e)}")
        return None

if __name__ == "__main__":
    # 设置参数
    results_directory = "results1"  # 可以根据需要修改路径
    output_filename = "ionization_data.npy"  # 输出文件名
    
    # 执行数据提取
    extract_seventh_row_data(results_directory, output_filename)
    
    
    print(f"\n脚本执行完成！")
    print(f"数据已保存到: {output_filename}")
    print(f"你可以使用 np.load('{output_filename}') 来加载数据")
